using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;

namespace ExcelRefresh
{
    public sealed class StaTaskScheduler : TaskScheduler, IDisposable
    {
        private readonly List<Thread> _threads;
        private BlockingCollection<Task> _tasks;
        
        public StaTaskScheduler(int numberOfThreads = 0)
        {
            if (numberOfThreads < 1)
                numberOfThreads = Environment.ProcessorCount;

            _tasks = new BlockingCollection<Task>();
            _threads = Enumerable.Range(0, numberOfThreads).Select(InitializeThread).ToList();
            _threads.ForEach(t => t.Start());
        }

        private Thread InitializeThread(int i)
        {
            var thread = new Thread(OnStart) {IsBackground = false};
            thread.SetApartmentState(ApartmentState.STA);
            return thread;
        }

        private void OnStart()
        {
            foreach (var t in _tasks.GetConsumingEnumerable())
                TryExecuteTask(t);
        }

        protected override void QueueTask(Task task)
        {
            _tasks.Add(task);
        }

        protected override IEnumerable<Task> GetScheduledTasks()
        {
            return _tasks.ToArray();
        }

        protected override bool TryExecuteTaskInline(
            Task task, 
            bool taskWasPreviouslyQueued)
        {
            return Thread.CurrentThread.GetApartmentState() == ApartmentState.STA && TryExecuteTask(task);
        }

        public override int MaximumConcurrencyLevel
        {
            get { return _threads.Count; }
        }

        public void Dispose()
        {
            if (_tasks != null)
            {
                _tasks.CompleteAdding();

                foreach (var thread in _threads) thread.Join();

                _tasks.Dispose();
                _tasks = null;
            }
        }
    }
}